"""这是一个多模态检索增强生成（RAG）系统，实现以下完整流程：
存储：将图片+文本描述存入向量数据库（ChromaDB）
检索：用文本查询检索相似图片
理解：用GLM-4V模型分析图片内容
生成：用CogView-3模型根据描述生成新图片

"""

import chromadb
from chromadb.utils import embedding_functions
import requests
from pathlib import Path
import base64
from PIL import Image
import io
import os
from dotenv import load_dotenv

load_dotenv() # 加载.env文件中的环境变量

# 1. 初始化ChromaDB客户端
def initialize_chroma_client():
    chroma_client = chromadb.PersistentClient(path="multimodal_db")
    return chroma_client

# 2. 创建多模态集合
"""核心逻辑：
使用本地部署的bge-m3模型生成嵌入向量（替代OpenAI API）
设置余弦相似度作为距离度量
hnsw:space: 指定向量索引类型（HNSW图算法）
"""
def create_multimodal_collection(chroma_client):
    # 使用OpenAI的多模态嵌入模型
    # multimodal_ef = embedding_functions.OpenAIEmbeddingFunction(
    #     api_key=os.getenv("OPENAI_API_KEY"),
    #     model_name="text-embedding-3-large"
    # )
    multimodal_ef = embedding_functions.SentenceTransformerEmbeddingFunction(r"D:\ideaSpace\MyPython\models\bge-m3")

    collection = chroma_client.create_collection(
        name="animals",
        embedding_function=multimodal_ef,
        metadata={"hnsw:space": "cosine"}
    )
    print("多模态集合 'animals' 创建完成")
    return collection

# 3. 插入多模态数据（图片）
"""数据处理流程：
    1.图片→Base64编码
    2.生成简单文本描述
    3.将文本+元数据存入集合
  设计要点：
    使用mediaType字段支持未来扩展其他媒体类型"""
def insert_multimodal_data(collection, image_path, name="puppy"):
    # 读取图片并转换为base64
    with open(image_path, "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')

    # 生成图片描述作为文本嵌入的参考
    description = f"A photo of a {name}"

    collection.add(
        documents=[description],
        metadatas=[{"name": name, "mediaType": "image", "image_base64": image_base64}],
        ids=[f"img_{name}"]
    )
    print(f"图片数据 '{name}' 已插入")

# 4. 检索图片（基于文本查询）
"""检索逻辑：
基于文本查询找到最相似的图片
where参数过滤只返回图片类型数据
"""
def retrieve_image(collection, query):
    results = collection.query(
        query_texts=[query],
        n_results=1,
        where={"mediaType": "image"},
        include=["metadatas", "documents"]
    )

    if not results['ids'][0]:
        print("未找到匹配的图片")
        return None

    result = results['metadatas'][0][0]
    print("检索到的图片信息:", result['name'])
    return result['image_base64']

# 5. 用glm-4v生成图片描述
def generate_description_from_image_zhipu(prompt, image64, zhipu_api_key):
    """
    用智谱 glm-4v 实现“图片→描述”
    :param prompt: 用户提示
    :param image64: base64 编码的图片
    :param zhipu_api_key: 智谱控制台生成的 API Key
    :return: 描述文本
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {zhipu_api_key}"
    }

    payload = {
        "model": "glm-4v",          # 带视觉能力的模型
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image64}"}}
                ]
            }
        ],
        "max_tokens": 300
    }

    resp = requests.post(
        "https://open.bigmodel.cn/api/paas/v4/chat/completions",
        headers=headers,
        json=payload,
        timeout=30
    )
    resp.raise_for_status()
    return resp.json()['choices'][0]['message']['content']

# 6. 用cogview-3根据描述生成图片
def generate_image_zhipuai(prompt: str, zhipuai_api_key: str, save_path: str = "generated_image.png") -> str:
    """
    使用智谱AI生成图片（最终修正版）
    :param prompt: 图片描述文本
    :param zhipuai_api_key: 有效的API Key
    :param save_path: 图片保存路径
    :return: 图片本地路径或None（失败时）
    """
    # 确保保存目录存在
    Path(save_path).parent.mkdir(parents=True, exist_ok=True)

    # API请求配置
    url = "https://open.bigmodel.cn/api/paas/v4/images/generations"
    headers = {
        "Authorization": f"Bearer {zhipuai_api_key}",
        "Content-Type": "application/json"
    }
    """关键参数：
        size: 控制生成图片分辨率
        response_format: 选择返回URL而非Base64（降低带宽占用）"""
    payload = {
        "model": "cogview-3",
        "prompt": prompt,
        "n": 1,
        "size": "1024x1024",
        "response_format": "url"
    }

    try:
        # 1. 创建图片生成任务
        response = requests.post(url, headers=headers, json=payload, timeout=30)
        response.raise_for_status()
        response_data = response.json()

        # 2. 检查返回数据结构
        if not response_data.get("data") or not isinstance(response_data["data"], list):
            raise ValueError("API返回数据格式异常")

        # 3. 获取图片URL
        image_url = response_data["data"][0]["url"]
        print(f"获取到图片URL: {image_url}")

        # 4. 下载图片
        image_response = requests.get(image_url, timeout=30)
        image_response.raise_for_status()

        # 5. 保存图片
        with open(save_path, "wb") as f:
            f.write(image_response.content)

        print(f"图片已成功保存到: {save_path}")
        return save_path

    except Exception as e:
        print(f"发生错误: {str(e)}")
        print("完整API响应:", response_data if 'response_data' in locals() else "无响应数据")
        return None

# 辅助函数：显示base64图片
def display_image_from_base64(image_base64):
    image_data = base64.b64decode(image_base64)
    image = Image.open(io.BytesIO(image_data))
    image.show()

if __name__ == "__main__":
    # 主流程控制：初始化DB → 插入数据 → 检索 → 生成描述 → 生成新图片
    # 1. 初始化ChromaDB客户端
    chroma_client = initialize_chroma_client()

    try:
        # 2. 创建集合
        collection = create_multimodal_collection(chroma_client)
    except:
        # 如果集合已存在，则获取它
        collection = chroma_client.get_collection(
            name="animals",
            # embedding_function=embedding_functions.OpenAIEmbeddingFunction(
            #     api_key=OPENAI_API_KEY,
            #     model_name="text-embedding-3-large"
            # )
            embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(r"D:\ideaSpace\MyPython\models\bge-m3")
        )

    # 3. 插入数据（替换为你的图片路径）
    image_path = "./data/astronaut_dog.png"  # 替换为你的图片路径
    if os.path.exists(image_path):
        insert_multimodal_data(collection, image_path, name="cute_puppy")
    else:
        print(f"警告: 图片文件 {image_path} 不存在，跳过数据插入步骤")

    # 4. 检索图片
    query_text = "a cute dog"
    retrieved_image = retrieve_image(collection, query_text)

    if retrieved_image:
        # 显示检索到的图片
        display_image_from_base64(retrieved_image)

        # 5. 用智谱 glm-4v 生成描述
        API_KEY = os.getenv("ZHIPUAI_API_KEY")
        description = generate_description_from_image_zhipu(
            prompt="这是一张我的宠物的图片，请给出可爱生动的描述，适合用于生成新的艺术风格图片。",
            image64=retrieved_image,
            zhipu_api_key=API_KEY,
        )

        # 6. 用智谱 cogview-3 生成新图片
        generate_image_zhipuai(
            prompt=f"根据以下描述创作一幅精美的数字艺术作品: {description}",
            zhipuai_api_key=API_KEY,
            save_path="./data/generated_artwork_dog.png"
        )
    else:
        print("无法继续处理，因为没有检索到图片")